# PBAR_Zspec_ExamineMCSpectra.py
# import PBAR_Zspec

import numpy as np
import matplotlib.pyplot as plt 

# list of bin edges
binEdgesList = [np.arange(256), \
                (np.arange(33)*3 + 7)\
                ]
binCentersList = []
for i in xrange(len(binEdgesList)):
    binCentersList.append((binEdgesList[i][1:] + binEdgesList[i][0:-1])/2.0)

binSizeList = np.array([1, 3])

outputDir = 'C:\Users\jkwong\Documents\Work\PBAR\dataMC'

prefix = 'dy04'

fullFileNameList = []
fileNameList = []
fileNameList.append(prefix + '_b0_255_1_t1_T0.csv')
fileNameList.append(prefix + '_b0_255_1_t5_T0.csv')

fileNameList.append(prefix + '_b7_103_3_t1_T0.csv')
fileNameList.append(prefix + '_b7_103_3_t5_T0.csv')

for i in xrange(len(fileNameList)):
    fullFileNameList.append(outputDir + '\\' + fileNameList[i])

fullFileNameList = np.array(fullFileNameList)

binTypeList = np.array([0, 0, 1, 1])
acquisitionTimeList = np.array([1, 5, 1, 5])

dat = []
for ii in xrange(len(fullFileNameList)):
    temp = np.genfromtxt(fullFileNameList[ii], \
                            delimiter=',', \
                            skip_header = 1, \
                            skip_footer = 0, \
                            dtype = 'uint32')
    dat.append(temp)


# calculate statistics

thresholdBin = 7
stats = {}

# Initialize array
statsList = ['binMean', 'binMeanSq', 'binSum', 'binSTD']
for ii in range(0,len(statsList)):
    stats[statsList[ii]] = np.zeros((len(dat), dat[ii].shape[1]))

for ii in xrange(len(dat)):
    binArray = np.matlib.repmat(binCentersList[binTypeList[ii]], dat[ii].shape[1],1).T
    
    datTemp = dat[ii].astype(float)        
    
    # remove the lower bins by setting it to zero
    cut = binCentersList[binTypeList[ii]] <= 7
    datTemp[cut,:] = 0

    stats['binMean'][ii,:] = (binArray * datTemp).sum(axis = 0).astype(float) / datTemp.sum(axis = 0).astype(float)
    stats['binMeanSq'][ii,:] = ((binArray**2) * datTemp).sum(axis = 0).astype(float) / datTemp.sum(axis = 0).astype(float)
    stats['binSum'][ii,:] = datTemp.sum(axis = 0).astype(float)
    
#    stats['binMean_q50'][ii,:] = stats['binMean'][ii,:] - stats['q_50'][ii,:]
    
    temp = np.matlib.repmat(stats['binMean'][ii,:], datTemp.shape[0],1)  # array of means, 256 x number detectors
    
    stats['binSTD'][ii,:] = np.sqrt(  ( ((binArray - temp)**2) * datTemp).sum(axis = 0).astype(float) / datTemp.sum(axis = 0).astype(float)  )



# plot all spectra
plt.figure()
plt.grid()

detNumber = 70

for ii in xrange(len(dat)):
    normalization = 1.0/acquisitionTimeList[ii] / binSizeList[binTypeList[ii]]
    plot(binCentersList[binTypeList[ii]], \
    dat[ii][:,detNumber]*normalization, \
    label = ('%s, t = %d, bin size = %d, bin mean = %f, bin std = %f, normalization = %f' \
    %(fileNameList[ii], acquisitionTimeList[ii], binSizeList[binTypeList[ii]], \
    stats['binMean'][ii, detNumber], stats['binSTD'][ii, detNumber], normalization)), linewidth = 2)
plt.xlim((0, 200))
plt.yscale('log')
plt.ylabel('Counts')
plt.xlabel('Bin')
plt.legend()
plt.show()